package com.linkhub.oms.component.trace.util;

import com.linkhub.oms.component.trace.core.QtTransferInfo;
import org.apache.commons.lang3.StringUtils;

import java.util.function.Supplier;

import static com.linkhub.oms.common.constant.TraceConstant.BASE_SPAN_ID_CONCAT_SPAN_ID;
import static com.linkhub.oms.common.constant.TraceConstant.REQUEST_ID_CONCAT_SPAN_ID;


/**
 * @Description: TraceInfoAssist
 * @Author: echo
 * @Date: 2022/9/1 18:04
 * @Version: 1.0
 */
public class TraceInfoAssist {

    public static String getTracePrefixKey(QtTransferInfo qtTransferInfo) {
        return String.format("%s%s%s", qtTransferInfo.getRequestId(),
                REQUEST_ID_CONCAT_SPAN_ID,
                qtTransferInfo.getParentSpanId());
    }

    public static String getTraceId(QtTransferInfo qtTransferInfo, int spanId) {
        String tracePrefixKey = getTracePrefixKey(qtTransferInfo);
        spanId = spanId == 0 ? qtTransferInfo.getLeafSpanIdInfo().get(tracePrefixKey) : spanId;
        return String.format("%s%s%s", tracePrefixKey,
                BASE_SPAN_ID_CONCAT_SPAN_ID,
                spanId);
    }

    public static String getTraceId(String requestId, String parentSpanId) {
        return String.format("%s%s%s", requestId,
                REQUEST_ID_CONCAT_SPAN_ID,
                parentSpanId);
    }


    public static String getSpanId(QtTransferInfo qtTransferInfo) {
        String tracePrefixKey = getTracePrefixKey(qtTransferInfo);
        return String.format("%s%s%s", qtTransferInfo.getParentSpanId(),
                BASE_SPAN_ID_CONCAT_SPAN_ID,
                qtTransferInfo.getLeafSpanIdInfo().get(tracePrefixKey));
    }

    public static String getTraceInfoIfNotBlank(Supplier<String> traceIdGetter) {
        if (traceIdGetter == null) {
            return null;
        }

        String traceId = traceIdGetter.get();

        return StringUtils.defaultIfBlank(traceId, StringUtils.EMPTY);

    }
}
